import os
from enum import Enum

import requests
from langchain.agents import AgentExecutor
from langchain.agents import tool
from langchain.chat_models import QianfanChatEndpoint
from langchain.pydantic_v1 import BaseModel, Field
from qianfan.extensions.langchain.agents import QianfanSingleActionAgent


os.environ["QIANFAN_AK"] = "KBe3HbStsjvnGS9KFrNZFPes"
os.environ["QIANFAN_SK"] = "3yGoXGIdCCy8QLwcLCLdh3UGGlM7PKNn"

class TemperatureUnitEnum(str, Enum):
    celsius: str = "摄氏度"
    fahrenheit: str = "华氏度"


class WeatherToolSchema(BaseModel):
    """获得指定地点的天气"""

    location: str = Field(description="省，市名，例如：河北省，石家庄")
    unit: TemperatureUnitEnum = Field(description="温度单位")


@tool(args_schema=WeatherToolSchema)
def get_current_weather(location: str, unit: TemperatureUnitEnum) -> str:
    """获得指定地点的天气"""

    print("\n" + location)
    print("\n" + unit)

    geo_url = "https://restapi.amap.com/v3/geocode/geo"
    geo_params = {'key': '58b8527d0f9ed3264778f088c348ccb6', 'address': location, 'output': 'JSON'}
    ad_code = requests.get(geo_url, params=geo_params).json().get("geocodes")[0].get("adcode")
    print("\n城市：" + location + ", 编码：" + ad_code)

    weather_url = "https://restapi.amap.com/v3/weather/weatherInfo"
    weather_params = {'key': '58b8527d0f9ed3264778f088c348ccb6', 'city': ad_code, 'extensions': 'base', 'output': 'JSON'}
    lives_ = requests.get(weather_url, params=weather_params).json().get("lives")[0]
    temperature = lives_.get("temperature")
    return temperature


tools = [get_current_weather]
# 目前千帆仅EB模型支持自定义agent与工具、EB-4后续会支持
llm = QianfanChatEndpoint(model="ERNIE-Bot")
# QianfanSingleActionAgent 为串行执行，一次规划只能使用一个工具；
agent = QianfanSingleActionAgent.from_system_prompt(tools, llm)
agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)

while True:
    input_str = input("请输入您的问题：")
    # result = agent_executor.run("今天上海的气温是多少摄氏度")
    result = agent_executor.run(input_str)
    print("result: \n" + result)
